import numpy as np
import pickle

temp = [100]*100


def qrs_detect(temp):
    stage = "Q"
    wave_yuzhi = 18000
    q_candi = 100000
    q_score = 0
    q_yuzhi = 10000
    r_candi = 0
    r_score = 0
    r_yuzhi = 10000
    s_candi = 100000
    s_score = 0
    s_yuzhi = 13

    q_list, r_list, s_list = [], [], []

    for idx, item in enumerate(temp):
        if stage == "Q":
            if item < q_candi:
                q_candi, q_candi_idx = item, idx
                q_score = 0
            else:
                q_score += (item - q_candi)
            if q_score > q_yuzhi:
                r_candi, r_candi_idx = item, idx
                q_score = 0
                r_score = 0
                stage = "R"
                continue
        if stage == "R":
            if item > r_candi:
                r_candi, r_candi_idx = item, idx
            else:
                r_score += (r_candi - item)
            if r_score > r_yuzhi:
                if r_candi - q_candi > wave_yuzhi:
                    s_candi, s_candi_idx = item, idx
                    r_score = 0
                    stage = "S"
                    continue
                else:
                    stage = "Q"
                    continue
        if stage == "S":
            if item < s_candi:
                s_candi, s_candi_idx = item, idx
            else:
                s_score += 1
            if s_score > s_yuzhi:
                s_score = 0
                q_list.append(q_candi_idx)
                r_list.append(r_candi_idx)
                s_list.append(s_candi_idx)

                print(q_candi_idx, r_candi_idx, s_candi_idx)
                q_candi = 100000
                r_candi = 0
                s_candi = 100000
                stage = "Q"
                continue
    return q_list, r_list, s_list, stage


with open("data_collected/test23.p", "rb") as f:
    wave = pickle.load(f)
q_list, r_list, s_list, stage = qrs_detect(wave)
print(stage)
import matplotlib.pyplot as plt
wave = np.array(wave)
q_x, r_x, s_x = np.array(q_list), np.array(r_list), np.array(s_list)
q_y, r_y, s_y = wave[q_x], wave[r_x], wave[s_x]
plt.plot(wave)
plt.scatter(q_x, q_y)
plt.scatter(r_x, r_y)
plt.scatter(s_x, s_y)
plt.show()
